#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
UMA/UKM Metrics Extractor

This script extracts UMA (User Metrics Analytics) and UKM (User Key Metrics)
definitions from Chromium codebase and stores them in a SQLite database.

Copyright 2024 Chromium Authors
Use of this source code is governed by a BSD-style license that can be
found in the LICENSE file.
"""

import argparse
import concurrent.futures
import logging
import os
import sqlite3
import sys
import time
import xml.etree.ElementTree as ET
from dataclasses import dataclass
from pathlib import Path
from typing import Dict, List, Optional, Set, Tuple, Union


# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(levelname)s - %(message)s',
    handlers=[
        logging.StreamHandler(sys.stdout),
        logging.FileHandler('uma_ukm_extractor.log', encoding='utf-8')
    ]
)
logger = logging.getLogger(__name__)


@dataclass
class MetricInfo:
    """Base class for metric information."""
    name: str
    type: Optional[str] = None
    summary: Optional[str] = None
    owners: List[str] = None
    file_path: str = ""
    line_number: int = 0

    def __post_init__(self):
        if self.owners is None:
            self.owners = []
        if not self.type:
            self.type = "unknown"


@dataclass
class HistogramInfo(MetricInfo):
    """Histogram metric information."""
    units: Optional[str] = None
    expires_after: Optional[str] = None
    component: Optional[str] = None
    enum_name: Optional[str] = None
    improvement_direction: Optional[str] = None
    deprecated: bool = False
    obsolete_message: Optional[str] = None
    tokens: Dict[str, str] = None

    def __post_init__(self):
        super().__post_init__()
        self.type = "histogram"
        if self.tokens is None:
            self.tokens = {}


@dataclass
class EnumInfo(MetricInfo):
    """Enumeration information."""
    type_name: Optional[str] = None
    values: List[Tuple[int, str, Optional[str]]] = None

    def __post_init__(self):
        super().__post_init__()
        self.type = "enum"
        if self.values is None:
            self.values = []


@dataclass
class UKMEventInfo(MetricInfo):
    """UKM Event information."""
    singular: bool = False
    metrics: List[Tuple[str, Optional[str], bool, Optional[str]]] = None  # (name, summary, deprecated, obsolete)

    def __post_init__(self):
        super().__post_init__()
        self.type = "ukm_event"
        if self.metrics is None:
            self.metrics = []


class FileScanner:
    """Scanner for discovering UMA/UKM definition files."""

    def __init__(self, chromium_dir: str):
        self.chromium_dir = Path(chromium_dir)
        self.histogram_files = []
        self.enum_files = []
        self.ukm_files = []

    def scan(self) -> Tuple[List[Path], List[Path], List[Path]]:
        """Scan directory for UMA/UKM XML files."""
        logger.info(f"Scanning directory: {self.chromium_dir}")

        # Patterns for different file types
        histogram_patterns = [
            "**/histograms.xml",
            "**/metadata/*/histograms.xml"
        ]

        enum_patterns = [
            "**/enums.xml",
            "**/metadata/*/enums.xml"
        ]

        ukm_patterns = [
            "**/ukm.xml"
        ]

        # Scan for files
        for pattern in histogram_patterns:
            self.histogram_files.extend(self.chromium_dir.glob(pattern))

        for pattern in enum_patterns:
            self.enum_files.extend(self.chromium_dir.glob(pattern))

        for pattern in ukm_patterns:
            self.ukm_files.extend(self.chromium_dir.glob(pattern))

        # Remove duplicates and sort
        self.histogram_files = sorted(list(set(self.histogram_files)))
        self.enum_files = sorted(list(set(self.enum_files)))
        self.ukm_files = sorted(list(set(self.ukm_files)))

        logger.info(f"Found {len(self.histogram_files)} histogram files")
        logger.info(f"Found {len(self.enum_files)} enum files")
        logger.info(f"Found {len(self.ukm_files)} UKM files")

        return self.histogram_files, self.enum_files, self.ukm_files


class XMLParser:
    """Parser for extracting metrics from XML files."""

    def __init__(self):
        self.namespace = {'': ''}  # Default namespace
        self.global_variants = {}  # Store global variant definitions

    def parse_histogram_file(self, file_path: Path) -> List[HistogramInfo]:
        """Parse histograms XML file."""
        histograms = []

        try:
            tree = ET.parse(file_path)
            root = tree.getroot()

            # First, extract global variants definitions
            self._extract_global_variants(root)

            # Then extract histograms
            for hist_elem in root.findall('.//histogram'):
                hist_info = self._extract_histogram_info(hist_elem, file_path)
                if hist_info:
                    histograms.append(hist_info)

        except ET.ParseError as e:
            logger.error(f"XML parsing error in {file_path}: {e}")
        except Exception as e:
            logger.error(f"Error parsing histogram file {file_path}: {e}")

        return histograms

    def _extract_global_variants(self, root: ET.Element):
        """Extract global variants definitions from the XML."""
        try:
            for variants_elem in root.findall('.//variants'):
                variants_name = variants_elem.get('name')
                if not variants_name:
                    continue

                variants = []
                for variant_elem in variants_elem.findall('variant'):
                    variant_name = variant_elem.get('name')
                    variant_summary = variant_elem.get('summary')
                    if variant_name:
                        variants.append((variant_name, variant_summary))

                if variants:
                    self.global_variants[variants_name] = variants

        except Exception as e:
            logger.debug(f"Error extracting global variants: {e}")

    def parse_enum_file(self, file_path: Path) -> List[EnumInfo]:
        """Parse enums XML file."""
        enums = []

        try:
            tree = ET.parse(file_path)
            root = tree.getroot()

            for enum_elem in root.findall('.//enum'):
                enum_info = self._extract_enum_info(enum_elem, file_path)
                if enum_info:
                    enums.append(enum_info)

        except ET.ParseError as e:
            logger.error(f"XML parsing error in {file_path}: {e}")
        except Exception as e:
            logger.error(f"Error parsing enum file {file_path}: {e}")

        return enums

    def parse_ukm_file(self, file_path: Path) -> List[UKMEventInfo]:
        """Parse UKM XML file."""
        ukm_events = []

        try:
            tree = ET.parse(file_path)
            root = tree.getroot()

            for event_elem in root.findall('.//event'):
                event_info = self._extract_ukm_event_info(event_elem, file_path)
                if event_info:
                    ukm_events.append(event_info)

        except ET.ParseError as e:
            logger.error(f"XML parsing error in {file_path}: {e}")
        except Exception as e:
            logger.error(f"Error parsing UKM file {file_path}: {e}")

        return ukm_events

    def _extract_histogram_info(self, elem: ET.Element, file_path: Path) -> Optional[HistogramInfo]:
        """Extract histogram information from XML element."""
        try:
            name = elem.get('name')
            if not name:
                return None

            hist_info = HistogramInfo(
                name=name,
                file_path=str(file_path),
                line_number=elem.sourceline if hasattr(elem, 'sourceline') else 0
            )

            # Extract attributes
            hist_info.units = elem.get('units')
            hist_info.expires_after = elem.get('expires_after')
            hist_info.enum_name = elem.get('enum')

            # Extract child elements
            hist_info.summary = self._get_text(elem, 'summary')
            hist_info.component = self._get_text(elem, 'component')

            # Extract owners
            for owner_elem in elem.findall('owner'):
                owner_text = self._get_element_text(owner_elem)
                if owner_text:
                    hist_info.owners.append(owner_text.strip())

            # Extract improvement direction
            improvement_elem = elem.find('improvement')
            if improvement_elem is not None:
                hist_info.improvement_direction = improvement_elem.get('direction')

            # Extract obsolete information
            obsolete_elem = elem.find('obsolete')
            if obsolete_elem is not None:
                hist_info.deprecated = True
                hist_info.obsolete_message = self._get_element_text(obsolete_elem)

            # Extract tokens and variants
            for token_elem in elem.findall('token'):
                token_key = token_elem.get('key')
                if not token_key:
                    continue

                # Extract variants for this token
                variants = []
                for variant_elem in token_elem.findall('variant'):
                    variant_name = variant_elem.get('name')
                    variant_summary = variant_elem.get('summary')
                    if variant_name:
                        variants.append((variant_name, variant_summary))

                # Get variants attribute if present (backward compatibility)
                variants_attr = token_elem.get('variants')

                if variants:
                    # Store token with variants
                    hist_info.tokens[token_key] = {
                        'name': token_key,
                        'variants': variants
                    }
                elif variants_attr:
                    # Check if variants_attr refers to a global variants definition
                    if variants_attr in self.global_variants:
                        # Use global variants
                        hist_info.tokens[token_key] = {
                            'name': token_key,
                            'variants': self.global_variants[variants_attr]
                        }
                    else:
                        # Old format: simple string reference
                        hist_info.tokens[token_key] = variants_attr

            return hist_info

        except Exception as e:
            logger.error(f"Error extracting histogram info: {e}")
            return None

    def _extract_enum_info(self, elem: ET.Element, file_path: Path) -> Optional[EnumInfo]:
        """Extract enum information from XML element."""
        try:
            name = elem.get('name')
            if not name:
                return None

            enum_info = EnumInfo(
                name=name,
                file_path=str(file_path),
                line_number=elem.sourceline if hasattr(elem, 'sourceline') else 0
            )

            # Extract summary
            enum_info.summary = self._get_text(elem, 'summary')

            # Extract owners
            for owner_elem in elem.findall('owner'):
                owner_text = self._get_element_text(owner_elem)
                if owner_text:
                    enum_info.owners.append(owner_text.strip())

            # Extract enum values
            for int_elem in elem.findall('int'):
                value_str = int_elem.get('value')
                label = int_elem.get('label')

                if value_str and label:
                    try:
                        value = int(value_str)
                        summary = self._get_element_text(int_elem)
                        enum_info.values.append((value, label, summary))
                    except ValueError:
                        logger.warning(f"Invalid enum value: {value_str}")

            return enum_info

        except Exception as e:
            logger.error(f"Error extracting enum info: {e}")
            return None

    def _extract_ukm_event_info(self, elem: ET.Element, file_path: Path) -> Optional[UKMEventInfo]:
        """Extract UKM event information from XML element."""
        try:
            name = elem.get('name')
            if not name:
                return None

            event_info = UKMEventInfo(
                name=name,
                file_path=str(file_path),
                line_number=elem.sourceline if hasattr(elem, 'sourceline') else 0
            )

            # Extract attributes
            event_info.singular = elem.get('singular') == 'True'

            # Extract summary
            event_info.summary = self._get_text(elem, 'summary')

            # Extract owners
            for owner_elem in elem.findall('owner'):
                owner_text = self._get_element_text(owner_elem)
                if owner_text:
                    event_info.owners.append(owner_text.strip())

            # Extract obsolete information
            obsolete_elem = elem.find('obsolete')
            if obsolete_elem is not None:
                event_info.deprecated = True
                event_info.obsolete_message = self._get_element_text(obsolete_elem)

            # Extract metrics
            for metric_elem in elem.findall('metric'):
                metric_name = metric_elem.get('name')
                if metric_name:
                    metric_summary = self._get_text(metric_elem, 'summary')

                    metric_deprecated = False
                    metric_obsolete = None

                    metric_obsolete_elem = metric_elem.find('obsolete')
                    if metric_obsolete_elem is not None:
                        metric_deprecated = True
                        metric_obsolete = self._get_element_text(metric_obsolete_elem)

                    event_info.metrics.append((metric_name, metric_summary, metric_deprecated, metric_obsolete))

            return event_info

        except Exception as e:
            logger.error(f"Error extracting UKM event info: {e}")
            return None

    def _get_text(self, elem: ET.Element, tag: str) -> Optional[str]:
        """Get text content of a child element."""
        child = elem.find(tag)
        return self._get_element_text(child) if child is not None else None

    def _get_element_text(self, elem: ET.Element) -> Optional[str]:
        """Get text content of an element."""
        if elem is None:
            return None

        # Check if element has direct text
        if elem.text:
            return elem.text.strip()

        # Check if element has child text nodes
        text_content = []
        for child in elem:
            if child.text:
                text_content.append(child.text.strip())
            if child.tail:
                text_content.append(child.tail.strip())

        if text_content:
            return ' '.join(text_content)

        return None


class DatabaseManager:
    """Manager for SQLite database operations."""

    def __init__(self, db_path: str):
        self.db_path = db_path
        self.conn = None

    def connect(self):
        """Connect to database."""
        self.conn = sqlite3.connect(self.db_path)
        self.conn.row_factory = sqlite3.Row  # Enable dictionary access

    def close(self):
        """Close database connection."""
        if self.conn:
            self.conn.close()

    def update_histogram_enum_relations(self):
        """Update enum_id for histograms that have enum_name but missing enum_id."""
        cursor = self.conn.cursor()

        logger.info("Updating histogram-enum relations...")

        # Find histograms with enum_name but missing enum_id
        cursor.execute("""
            SELECT metric_id, enum_name
            FROM histogram_details
            WHERE enum_name IS NOT NULL
            AND enum_name <> ''
            AND enum_id IS NULL
        """)

        histograms_to_update = cursor.fetchall()

        if not histograms_to_update:
            logger.info("No histogram-enum relations need updating")
            return

        logger.info(f"Found {len(histograms_to_update)} histograms to update")

        updated_count = 0
        for metric_id, enum_name in histograms_to_update:
            cursor.execute('SELECT id FROM metrics WHERE name = ? AND type = ?', (enum_name, 'enum'))
            enum_row = cursor.fetchone()

            if enum_row:
                enum_id = enum_row[0]
                cursor.execute("""
                    UPDATE histogram_details
                    SET enum_id = ?
                    WHERE metric_id = ?
                """, (enum_id, metric_id))
                updated_count += 1
                logger.debug(f"Updated histogram {metric_id} with enum_id {enum_id} for enum '{enum_name}'")
            else:
                logger.debug(f"No enum found for enum_name '{enum_name}'")

        self.conn.commit()
        logger.info(f"Successfully updated {updated_count} histogram-enum relations")

    def create_tables(self):
        """Create database tables."""
        cursor = self.conn.cursor()

        # Main metrics table
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS metrics (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                name TEXT NOT NULL,
                type TEXT NOT NULL CHECK(type IN ('histogram', 'enum', 'ukm_event', 'ukm_metric')),
                summary TEXT,
                units TEXT,
                expires_after TEXT,
                component TEXT,
                singular BOOLEAN DEFAULT FALSE,
                file_path TEXT NOT NULL,
                line_number INTEGER,
                created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                UNIQUE(name, type)
            )
        ''')

        # Metric owners table
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS metric_owners (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                metric_id INTEGER NOT NULL,
                owner_email TEXT NOT NULL,
                FOREIGN KEY (metric_id) REFERENCES metrics (id) ON DELETE CASCADE,
                UNIQUE(metric_id, owner_email)
            )
        ''')

        # Histogram details table
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS histogram_details (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                metric_id INTEGER NOT NULL UNIQUE,
                enum_name TEXT,
                enum_id INTEGER,
                improvement_direction TEXT CHECK(improvement_direction IN ('HIGHER_IS_BETTER', 'LOWER_IS_BETTER', 'NEITHER_IS_BETTER')),
                deprecated BOOLEAN DEFAULT FALSE,
                obsolete_message TEXT,
                FOREIGN KEY (metric_id) REFERENCES metrics (id) ON DELETE CASCADE,
                FOREIGN KEY (enum_id) REFERENCES metrics (id) ON DELETE SET NULL
            )
        ''')

        # Enum values table
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS enum_values (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                enum_id INTEGER NOT NULL,
                value INTEGER NOT NULL,
                label TEXT NOT NULL,
                summary TEXT,
                FOREIGN KEY (enum_id) REFERENCES metrics (id) ON DELETE CASCADE,
                UNIQUE(enum_id, value)
            )
        ''')

        # UKM metrics table
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS ukm_metrics (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                event_id INTEGER NOT NULL,
                metric_name TEXT NOT NULL,
                summary TEXT,
                deprecated BOOLEAN DEFAULT FALSE,
                obsolete_message TEXT,
                FOREIGN KEY (event_id) REFERENCES metrics (id) ON DELETE CASCADE,
                UNIQUE(event_id, metric_name)
            )
        ''')

        # Token and variant tables
        cursor.execute('''
            CREATE TABLE IF NOT EXISTS tokens (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                token_name TEXT NOT NULL UNIQUE
            )
        ''')

        cursor.execute('''
            CREATE TABLE IF NOT EXISTS variants (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                token_id INTEGER NOT NULL,
                variant_name TEXT NOT NULL,
                variant_summary TEXT,
                FOREIGN KEY (token_id) REFERENCES tokens (id) ON DELETE CASCADE,
                UNIQUE(token_id, variant_name)
            )
        ''')

        cursor.execute('''
            CREATE TABLE IF NOT EXISTS histogram_tokens (
                id INTEGER PRIMARY KEY AUTOINCREMENT,
                metric_id INTEGER NOT NULL,
                token_key TEXT NOT NULL,
                token_name TEXT NOT NULL,
                FOREIGN KEY (metric_id) REFERENCES metrics (id) ON DELETE CASCADE
            )
        ''')

        # Create indexes
        cursor.execute('CREATE INDEX IF NOT EXISTS idx_metrics_name ON metrics(name)')
        cursor.execute('CREATE INDEX IF NOT EXISTS idx_metrics_type ON metrics(type)')
        cursor.execute('CREATE INDEX IF NOT EXISTS idx_metrics_component ON metrics(component)')
        cursor.execute('CREATE INDEX IF NOT EXISTS idx_metric_owners_metric_id ON metric_owners(metric_id)')
        cursor.execute('CREATE INDEX IF NOT EXISTS idx_enum_values_enum_id ON enum_values(enum_id)')
        cursor.execute('CREATE INDEX IF NOT EXISTS idx_ukm_metrics_event_id ON ukm_metrics(event_id)')
        cursor.execute('CREATE INDEX IF NOT EXISTS idx_tokens_name ON tokens(token_name)')
        cursor.execute('CREATE INDEX IF NOT EXISTS idx_variants_token_id ON variants(token_id)')
        cursor.execute('CREATE INDEX IF NOT EXISTS idx_histogram_tokens_metric_id ON histogram_tokens(metric_id)')

        # New indexes for histogram-enum optimization
        cursor.execute('CREATE INDEX IF NOT EXISTS idx_histogram_details_enum_id ON histogram_details(enum_id)')
        cursor.execute('CREATE INDEX IF NOT EXISTS idx_histogram_details_metric_enum ON histogram_details(metric_id, enum_id)')

        self.conn.commit()
        logger.info("Database tables created successfully")

    def insert_histogram(self, hist_info: HistogramInfo) -> int:
        """Insert histogram information."""
        cursor = self.conn.cursor()

        # Check if metric already exists
        cursor.execute('SELECT id FROM metrics WHERE name = ? AND type = ?', (hist_info.name, hist_info.type))
        row = cursor.fetchone()

        if row:
            metric_id = row[0]
            # Update existing record
            cursor.execute('''
                UPDATE metrics
                SET summary = ?, units = ?, expires_after = ?, component = ?, 
                    file_path = ?, line_number = ?, updated_at = CURRENT_TIMESTAMP
                WHERE id = ?
            ''', (
                hist_info.summary, hist_info.units, hist_info.expires_after, 
                hist_info.component, hist_info.file_path, hist_info.line_number, 
                metric_id
            ))
        else:
            # Insert new record
            cursor.execute('''
                INSERT INTO metrics
                (name, type, summary, units, expires_after, component, file_path, line_number, updated_at)
                VALUES (?, ?, ?, ?, ?, ?, ?, ?, CURRENT_TIMESTAMP)
            ''', (
                hist_info.name, hist_info.type, hist_info.summary, hist_info.units,
                hist_info.expires_after, hist_info.component, hist_info.file_path, hist_info.line_number
            ))
            metric_id = cursor.lastrowid

        # Insert owners (always add new ones, ignore duplicates)
        for owner in hist_info.owners:
            cursor.execute('''
                INSERT OR IGNORE INTO metric_owners (metric_id, owner_email)
                VALUES (?, ?)
            ''', (metric_id, owner))

        # Find enum_id if enum_name is provided
        enum_id = None
        if hist_info.enum_name:
            cursor.execute('SELECT id FROM metrics WHERE name = ? AND type = ?', (hist_info.enum_name, 'enum'))
            enum_row = cursor.fetchone()
            if enum_row:
                enum_id = enum_row[0]
                logger.debug(f"Found enum_id {enum_id} for enum_name '{hist_info.enum_name}'")
            else:
                logger.debug(f"No enum found for enum_name '{hist_info.enum_name}'")

        # Insert histogram details with enum_id
        cursor.execute('''
            INSERT OR REPLACE INTO histogram_details
            (metric_id, enum_name, enum_id, improvement_direction, deprecated, obsolete_message)
            VALUES (?, ?, ?, ?, ?, ?)
        ''', (
            metric_id, hist_info.enum_name, enum_id, hist_info.improvement_direction,
            hist_info.deprecated, hist_info.obsolete_message
        ))

        # Insert token information
        for token_key, token_data in hist_info.tokens.items():
            if isinstance(token_data, dict):
                # New format: dict with variants
                token_name = token_data.get('name', token_key)
                variants = token_data.get('variants', [])

                # Insert token
                cursor.execute('''
                    INSERT OR IGNORE INTO tokens (token_name) VALUES (?)
                ''', (token_name,))

                cursor.execute('SELECT id FROM tokens WHERE token_name = ?', (token_name,))
                token_id = cursor.fetchone()[0]

                # Insert histogram-token relationship
                cursor.execute('''
                    INSERT OR REPLACE INTO histogram_tokens
                    (metric_id, token_key, token_name)
                    VALUES (?, ?, ?)
                ''', (metric_id, token_key, token_name))

                # Insert variants
                for variant_name, variant_summary in variants:
                    cursor.execute('''
                        INSERT OR REPLACE INTO variants
                        (token_id, variant_name, variant_summary)
                        VALUES (?, ?, ?)
                    ''', (token_id, variant_name, variant_summary))
            else:
                # Old format: simple string
                token_name = token_data
                cursor.execute('''
                    INSERT OR IGNORE INTO tokens (token_name) VALUES (?)
                ''', (token_name,))

                cursor.execute('SELECT id FROM tokens WHERE token_name = ?', (token_name,))
                token_id = cursor.fetchone()[0]

                # Insert histogram-token relationship
                cursor.execute('''
                    INSERT OR REPLACE INTO histogram_tokens
                    (metric_id, token_key, token_name)
                    VALUES (?, ?, ?)
                ''', (metric_id, token_key, token_name))

        self.conn.commit()
        return metric_id

    def insert_enum(self, enum_info: EnumInfo) -> int:
        """Insert enum information."""
        cursor = self.conn.cursor()

        # Check if metric already exists
        cursor.execute('SELECT id FROM metrics WHERE name = ? AND type = ?', (enum_info.name, enum_info.type))
        row = cursor.fetchone()

        if row:
            metric_id = row[0]
            # Update existing record
            cursor.execute('''
                UPDATE metrics
                SET summary = ?, file_path = ?, line_number = ?, updated_at = CURRENT_TIMESTAMP
                WHERE id = ?
            ''', (enum_info.summary, enum_info.file_path, enum_info.line_number, metric_id))
        else:
            # Insert new record
            cursor.execute('''
                INSERT INTO metrics
                (name, type, summary, file_path, line_number, updated_at)
                VALUES (?, ?, ?, ?, ?, CURRENT_TIMESTAMP)
            ''', (enum_info.name, enum_info.type, enum_info.summary, enum_info.file_path, enum_info.line_number))
            metric_id = cursor.lastrowid

        # Insert owners
        for owner in enum_info.owners:
            cursor.execute('''
                INSERT OR IGNORE INTO metric_owners (metric_id, owner_email)
                VALUES (?, ?)
            ''', (metric_id, owner))

        # Insert enum values
        for value, label, summary in enum_info.values:
            cursor.execute('''
                INSERT OR REPLACE INTO enum_values (enum_id, value, label, summary)
                VALUES (?, ?, ?, ?)
            ''', (metric_id, value, label, summary))

        self.conn.commit()
        return metric_id

    def insert_ukm_event(self, event_info: UKMEventInfo) -> int:
        """Insert UKM event information."""
        cursor = self.conn.cursor()

        # Check if metric already exists
        cursor.execute('SELECT id FROM metrics WHERE name = ? AND type = ?', (event_info.name, event_info.type))
        row = cursor.fetchone()

        if row:
            metric_id = row[0]
            # Update existing record
            cursor.execute('''
                UPDATE metrics
                SET summary = ?, singular = ?, file_path = ?, line_number = ?, updated_at = CURRENT_TIMESTAMP
                WHERE id = ?
            ''', (
                event_info.summary, event_info.singular, event_info.file_path, 
                event_info.line_number, metric_id
            ))
        else:
            # Insert new record
            cursor.execute('''
                INSERT INTO metrics
                (name, type, summary, singular, file_path, line_number, updated_at)
                VALUES (?, ?, ?, ?, ?, ?, CURRENT_TIMESTAMP)
            ''', (
                event_info.name, event_info.type, event_info.summary,
                event_info.singular, event_info.file_path, event_info.line_number
            ))
            metric_id = cursor.lastrowid

        # Insert owners
        for owner in event_info.owners:
            cursor.execute('''
                INSERT OR IGNORE INTO metric_owners (metric_id, owner_email)
                VALUES (?, ?)
            ''', (metric_id, owner))

        # Insert UKM metrics
        for metric_name, metric_summary, metric_deprecated, metric_obsolete in event_info.metrics:
            cursor.execute('''
                INSERT OR REPLACE INTO ukm_metrics
                (event_id, metric_name, summary, deprecated, obsolete_message)
                VALUES (?, ?, ?, ?, ?)
            ''', (metric_id, metric_name, metric_summary, metric_deprecated, metric_obsolete))

        self.conn.commit()
        return metric_id

    def get_stats(self) -> Dict[str, int]:
        """Get database statistics."""
        cursor = self.conn.cursor()

        stats = {}

        cursor.execute("SELECT COUNT(*) FROM metrics WHERE type = 'histogram'")
        stats['histograms'] = cursor.fetchone()[0]

        cursor.execute("SELECT COUNT(*) FROM metrics WHERE type = 'enum'")
        stats['enums'] = cursor.fetchone()[0]

        cursor.execute("SELECT COUNT(*) FROM metrics WHERE type = 'ukm_event'")
        stats['ukm_events'] = cursor.fetchone()[0]

        cursor.execute("SELECT COUNT(*) FROM ukm_metrics")
        stats['ukm_metrics'] = cursor.fetchone()[0]

        cursor.execute("SELECT COUNT(*) FROM enum_values")
        stats['enum_values'] = cursor.fetchone()[0]

        return stats


class MetricsExtractor:
    """Main extractor class."""

    def __init__(self, chromium_dir: str, db_path: str, workers: int = 4):
        self.scanner = FileScanner(chromium_dir)
        self.parser = XMLParser()
        self.db_manager = DatabaseManager(db_path)
        self.workers = workers
        # Store the base chromium directory for relative path conversion
        self.chromium_dir = Path(chromium_dir).resolve()

    def get_relative_path(self, file_path: str) -> str:
        """Convert absolute file path to relative path."""
        try:
            # Convert to Path object and make absolute
            abs_path = Path(file_path).resolve()

            # Get relative path from chromium directory
            rel_path = abs_path.relative_to(self.chromium_dir)

            # Convert to string with forward slashes for consistency
            return str(rel_path).replace(os.sep, '/')

        except ValueError:
            # If file is outside chromium directory, return original path
            # This shouldn't happen in normal usage, but handle gracefully
            logger.warning(f"File {file_path} is outside chromium directory, using original path")
            return file_path

    def run(self) -> Dict[str, int]:
        """Run the extraction process."""
        start_time = time.time()

        logger.info("Starting UMA/UKM metrics extraction")

        # Connect to database
        self.db_manager.connect()
        self.db_manager.create_tables()

        # Scan for files
        histogram_files, enum_files, ukm_files = self.scanner.scan()

        total_files = len(histogram_files) + len(enum_files) + len(ukm_files)
        processed_files = 0

        # Process files sequentially to avoid SQLite threading issues
        all_files = list(histogram_files) + list(enum_files) + list(ukm_files)

        for file_path in all_files:
            try:
                if file_path in histogram_files:
                    self._process_histogram_file(file_path)
                elif file_path in enum_files:
                    self._process_enum_file(file_path)
                elif file_path in ukm_files:
                    self._process_ukm_file(file_path)

                processed_files += 1
                if processed_files % 10 == 0:
                    logger.info(f"Processed {processed_files}/{total_files} files")

            except Exception as e:
                logger.error(f"Error processing file {file_path}: {e}")

        # Expand tokens for reverse lookup (default behavior)
        if hasattr(self, '_expand_tokens') and self._expand_tokens.__name__ != '<lambda>':
            logger.info("Expanding tokens for reverse lookup...")
            self._expand_tokens()
        else:
            logger.info("Token expansion skipped")

        # Update histogram-enum relations before closing database
        logger.info("Updating histogram-enum relations...")
        self.db_manager.update_histogram_enum_relations()

        # Get statistics
        stats = self.db_manager.get_stats()

        # Close database
        self.db_manager.close()

        elapsed_time = time.time() - start_time
        logger.info(f"Extraction completed in {elapsed_time:.2f} seconds")

        return stats

    def _expand_tokens(self):
        """Expand tokens for reverse lookup functionality."""
        try:
            from token_expander import TokenExpander
            expander = TokenExpander(self.db_manager.db_path)
            expander.connect()
            expander.create_expansion_tables()
            expander.expand_all_histograms()
            expander.close()
            logger.info("Token expansion completed successfully")
        except Exception as e:
            logger.warning(f"Token expansion failed: {e}")

    def _process_histogram_file(self, file_path: Path):
        """Process a single histogram file."""
        histograms = self.parser.parse_histogram_file(file_path)
        for hist in histograms:
            # Convert file_path to relative path before storing
            hist.file_path = self.get_relative_path(hist.file_path)
            self.db_manager.insert_histogram(hist)

        if histograms:
            logger.debug(f"Processed {len(histograms)} histograms from {self.get_relative_path(str(file_path))}")

    def _process_enum_file(self, file_path: Path):
        """Process a single enum file."""
        enums = self.parser.parse_enum_file(file_path)
        for enum in enums:
            # Convert file_path to relative path before storing
            enum.file_path = self.get_relative_path(enum.file_path)
            self.db_manager.insert_enum(enum)

        if enums:
            logger.debug(f"Processed {len(enums)} enums from {self.get_relative_path(str(file_path))}")

    def _process_ukm_file(self, file_path: Path):
        """Process a single UKM file."""
        ukm_events = self.parser.parse_ukm_file(file_path)
        for event in ukm_events:
            # Convert file_path to relative path before storing
            event.file_path = self.get_relative_path(event.file_path)
            self.db_manager.insert_ukm_event(event)

        if ukm_events:
            logger.debug(f"Processed {len(ukm_events)} UKM events from {self.get_relative_path(str(file_path))}")


def main():
    """Main function."""
    parser = argparse.ArgumentParser(description='Extract UMA/UKM metrics from Chromium codebase')
    parser.add_argument('--chromium-dir', default='../../../',
                        help='Chromium source directory (default: ../../../)')
    parser.add_argument('--output', default='../uma_insight.db',
                        help='Output database file (default: ../uma_insight.db)')
    parser.add_argument('--verbose', '-v', action='store_true',
                        help='Enable verbose logging')
    parser.add_argument('--workers', type=int, default=4,
                        help='Number of worker threads (default: 4)')
    parser.add_argument('--no-token-expansion', action='store_true',
                        help='Skip token expansion (faster but no reverse lookup)')

    args = parser.parse_args()

    # Configure logging level
    if args.verbose:
        logging.getLogger().setLevel(logging.DEBUG)

    # Validate chromium directory
    chromium_dir = os.path.abspath(args.chromium_dir)
    if not os.path.isdir(chromium_dir):
        logger.error(f"Chromium directory does not exist: {chromium_dir}")
        sys.exit(1)

    logger.info(f"Using chromium directory: {chromium_dir}")
    logger.info(f"Output database: {args.output}")
    if args.no_token_expansion:
        logger.info("Token expansion: DISABLED (faster processing)")
    else:
        logger.info("Token expansion: ENABLED (default)")

    # Run extractor
    extractor = MetricsExtractor(chromium_dir, args.output, args.workers)

    try:
        if args.no_token_expansion:
            # Override _expand_tokens method to skip expansion
            extractor._expand_tokens = lambda: None

        stats = extractor.run()

        # Print results
        print("\n" + "="*50)
        print("EXTRACTION RESULTS")
        print("="*50)
        print(f"Histograms: {stats['histograms']:,}")
        print(f"Enums: {stats['enums']:,}")
        print(f"UKM Events: {stats['ukm_events']:,}")
        print(f"UKM Metrics: {stats['ukm_metrics']:,}")
        print(f"Enum Values: {stats['enum_values']:,}")
        print(f"Database: {args.output}")
        if not args.no_token_expansion:
            print("Note: Tokens were expanded for reverse lookup functionality")
        print("="*50)

    except KeyboardInterrupt:
        logger.info("Extraction interrupted by user")
        sys.exit(1)
    except Exception as e:
        logger.error(f"Extraction failed: {e}")
        sys.exit(1)


if __name__ == '__main__':
    main()